package main

import (
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/websocket"
)

// 升级器 (upgrader) 配置。用于将 HTTP 连接升级为 WebSocket 连接。
var upgrader = websocket.Upgrader{
	// 允许跨域。在生产环境中，应设置为具体的、受信任的域名。
	CheckOrigin: func(r *http.Request) bool {
		return true // 允许所有来源。仅用于开发测试！
	},
}

// 处理 WebSocket 连接的函数
func wsHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("收到新的 WebSocket 连接请求...")

	// 1. 将 HTTP 连接升级为 WebSocket 连接
	//    `conn` 是升级后的 WebSocket 连接对象。
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket 升级失败: %v", err)
		return // Upgrade 方法已经处理了 HTTP 响应（如 400 Bad Request）
	}
	// 确保连接在函数退出时关闭
	defer conn.Close()

	fmt.Println("WebSocket 连接已建立！")

	// 2. 持续读取客户端发送的消息
	for {
		// ReadMessage 阻塞等待客户端消息
		messageType, message, err := conn.ReadMessage()
		if err != nil {
			// 如果是正常的连接关闭（如客户端断开），打印信息即可
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
				log.Printf("读取消息时发生意外错误: %v", err)
			} else {
				fmt.Println("客户端已断开连接")
			}
			break // 退出循环，结束此连接的处理
		}

		// 打印收到的消息（用于调试）
		fmt.Printf("收到消息 (类型: %d): %s\n", messageType, message)

		// 3. 将收到的消息原样回传给客户端（Echo 服务）
		//    你可以在这里添加任何业务逻辑，比如处理数据、广播给其他客户端等。
		if err := conn.WriteMessage(messageType, message); err != nil {
			log.Printf("发送消息失败: %v", err)
			break // 如果发送失败，通常意味着连接已断开，退出循环
		}
	}

	fmt.Println("WebSocket 连接已关闭。")
}

func main() {
	// 将路径 "/ws" 的请求映射到 wsHandler 处理函数
	http.HandleFunc("/ws", wsHandler)

	// 可选：提供一个简单的 HTML 页面用于测试（见下方调用方式）
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		html := `
        <html>
        <head>
            <title>WebSocket 测试</title>
        </head>
        <body>
            <h1>WebSocket Echo 测试</h1>
            <button onclick="connect()">连接 WebSocket</button>
            <button onclick="send()">发送 'Hello, Server!'</button>
            <button onclick="disconnect()">断开连接</button>
            <div id="output"></div>

            <script>
                let socket;
                const output = document.getElementById('output');

                function log(msg) {
                    const p = document.createElement('p');
                    p.textContent = msg;
                    output.appendChild(p);
                }

                function connect() {
                    // 创建 WebSocket 连接。确保 URL 正确（http -> ws）
                    socket = new WebSocket("ws://localhost:8080/ws");
                    
                    socket.onopen = function(event) {
                        log("连接已建立");
                    };

                    socket.onmessage = function(event) {
                        log("收到回复: " + event.data);
                    };

                    socket.onclose = function(event) {
                        log("连接已关闭");
                        socket = null;
                    };

                    socket.onerror = function(event) {
                        log("发生错误");
                    };
                }

                function send() {
                    if (socket && socket.readyState === WebSocket.OPEN) {
                        socket.send("Hello, Server! 这是来自浏览器的消息。");
                    } else {
                        log("WebSocket 未连接或未打开");
                    }
                }

                function disconnect() {
                    if (socket) {
                        socket.close();
                    }
                }
            </script>
        </body>
        </html>
        `
		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		fmt.Fprint(w, html)
	})

	fmt.Println("WebSocket 服务器正在启动...")
	fmt.Println("测试页面: http://localhost:8080")
	fmt.Println("WebSocket 端点: ws://localhost:8080/ws")
	fmt.Println("按 Ctrl+C 停止服务器")

	// 启动 HTTP 服务器，监听 8080 端口
	err := http.ListenAndServe(":8080", nil)
	if err != nil {
		log.Fatal("服务器启动失败: ", err)
	}
}
